package org.zachary.commonservice.controller;


import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import io.reactivex.Flowable;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpSession;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;

import java.io.IOException;
import java.io.PrintWriter;

import java.util.*;


@RestController
@RequestMapping("/chat")
public class ChatController {

    // 定义最大上下文长度
    private static final int MAX_HISTORY = 5;

    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public void stream(
            @RequestParam("prompt") String prompt,
            HttpServletRequest request,  // 添加request用于获取session
            HttpServletResponse response
    ) throws IOException {
        System.out.println("收到消息：" + prompt);

        HttpSession session = request.getSession(true);
        List<Message> history = (List<Message>) session.getAttribute("chatHistory");

        // 初始化对话历史
        if (history == null) {
            history = new ArrayList<>();
            session.setAttribute("chatHistory", history);
        }

        response.setContentType(MediaType.TEXT_EVENT_STREAM_VALUE);
        response.setCharacterEncoding("UTF-8");
        PrintWriter writer = response.getWriter();

        // 构建用户消息并添加到历史
        Message userMessage = Message.builder()
                .role(Role.USER.getValue())
                .content(prompt)
                .build();
        history.add(userMessage);

        // 控制历史长度
        if (history.size() > MAX_HISTORY * 2) {
            // 保留最新的 MAX_HISTORY 轮对话
            history = history.subList(history.size() - MAX_HISTORY * 2, history.size());
            session.setAttribute("chatHistory", history);
        }

        try {
            Generation generation = new Generation();
            StringBuilder assistantContent = new StringBuilder();

            // 流式接收响应
            Flowable<GenerationResult> stream = generation.streamCall(
                    GenerationParam.builder()
                            .apiKey("sk-2f79bca9b5cb42baaeb9b71d141b5672")
                            .model("qwen-plus")
                            .messages(history)  // 传入整个对话历史
                            .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                            .incrementalOutput(true)
                            .build()
            );

            stream.blockingForEach(result -> {
                String content = result.getOutput()
                        .getChoices()
                        .get(0)
                        .getMessage()
                        .getContent();
                if (content != null && !content.isEmpty()) {
                    writer.write("data: " + content + "\n\n");
                    writer.flush();
                    assistantContent.append(content);  // 累积AI回复内容
                }
            });

            // 将AI回复添加到历史
            history.add(Message.builder()
                    .role(Role.ASSISTANT.getValue())
                    .content(assistantContent.toString())
                    .build());

            session.setAttribute("chatHistory", history);  // 更新会话历史

            writer.write("data: [DONE]\n\n");
            writer.flush();

        } catch (Exception e) {
            // 出错时从历史中移除当前问题
            history.remove(history.size() - 1);
            writer.write("data: [ERROR] " + e.getMessage() + "\n\n");
            writer.flush();
        } finally {
            writer.close();
        }
    }

    // 添加清除上下文的接口
    @PostMapping("/clear")
    public ResponseEntity<String> clearContext(HttpServletRequest request) {
        HttpSession session = request.getSession(false);
        if (session != null) {
            session.removeAttribute("chatHistory");
        }
        return ResponseEntity.ok("上下文已清除");
    }
}


